'''
Author: caoniu caoniu@jushuitan.com
Date: 2023-06-07 20:20:46
LastEditors: caoniu caoniu@jushuitan.com
LastEditTime: 2023-06-21 04:15:53
FilePath: /py_project/chatWebsocket.py
Description: WebSocket请求
'''
import subprocess
import os
import time
import json
import asyncio
import uvicorn
from typing import Any, Dict, List

from langchain.agents import (
    initialize_agent,
    Tool,
    AgentType
)
from langchain.memory import (
    ConversationBufferMemory,
    ConversationSummaryBufferMemory
)
from langchain.prompts import (
    MessagesPlaceholder,
    MessagesPlaceholder
)
from langchain.schema import SystemMessage
from langchain.utilities import (
    GoogleSerperAPIWrapper
)
from langchain import (
    LLMMathChain,
)
from langchain.chat_models import ChatOpenAI
from langchain.callbacks import get_openai_callback

from fastapi import FastAPI, WebSocket, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

from bs4 import BeautifulSoup

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

# 配置环境变量
os.environ["OPENAI_API_KEY"] = "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"
os.environ["SERPER_API_KEY"] = "5745698f4a629390d7f489822b75f581439fee2a"
os.environ["SERPAPI_API_KEY"] = "b2ff07bc4b7494019705ef6b2f629300d9f6cc9366a92c28693dcd71062d0378"

class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str, websocket: WebSocket):
        for connection in self.active_connections:
            if connection is not websocket:
                await connection.send_text(message)


manager = ConnectionManager()


@app.get("/", response_class=HTMLResponse)
async def get(request: Request):
    return templates.TemplateResponse("chat.html", {"request": request})


#
# 生成llmAgent
#


def generateAgent():
    llm = ChatOpenAI(temperature=0,verbose=True, model="gpt-3.5-turbo-0613")
    # search = SerpAPIWrapper()
    search = GoogleSerperAPIWrapper()
    llm_math_chain = LLMMathChain.from_llm(llm=llm, verbose=True)
    tools = [
        Tool(
            name="Search",
            func=search.run,
            description="useful for when you need to answer questions or search something about current events. You should ask targeted questions",
        ),
        Tool(
            name="Calculator",
            func=llm_math_chain.run,
            description="useful for when you need to answer questions about math",
        ),
    ]

    agent_kwargs = {
        "extra_prompt_messages": [MessagesPlaceholder(variable_name="memory")],
        "system_message": SystemMessage(content="The following is a friendly conversation between a human and an AI. The AI is talkative and "
            "provides lots of specific details from its context. If the AI does not know the answer to a "
            "question, it truthfully says it does not know.")
    }
    memory = ConversationBufferMemory(
        memory_key="memory", return_messages=True)

    agent = initialize_agent(
        tools,
        llm,
        agent=AgentType.OPENAI_FUNCTIONS,
        verbose=True,
        agent_kwargs=agent_kwargs,
        memory=memory,
    )
    return agent


def getGoogleNewsData(prompt):
    search = GoogleSerperAPIWrapper(type="news")
    results = search.results(prompt)
    newsData = results["news"]
    return newsData


def getGoogleNewsHtml(newsData):
    # 构建HTML列表
    html = '<ul>'
    for item in newsData:
        html += '<li>标题: {}</li>'.format(item['title'])
        html += '<li>链接: <a href="{}">{}</a></li>'.format(
            item['link'], item['link'])
        html += '<li>简介: {}</li>'.format(item['snippet'])
        html += '<li>时间: {}</li>'.format(item['date'])
        html += '<li>来源: {}</li>'.format(item['source'])
        html += '<br/>'
    html += '</ul>'
    return html


def getGoogleNewsSummary(newsData):
    summary = ""
    for item in newsData:
        summary += str(item['position'])
        summary += "。"
        summary += item['title']
        summary += "。"
    return summary


# 移除html相关标签
def remove_html_tags(text):
    soup = BeautifulSoup(text, 'html.parser')
    return soup.get_text()


@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket)
    try:
        agent = generateAgent()
        await manager.send_personal_message(f"我是一个智能机器人助手，专门设计来回答你的问题和提供帮助。", websocket)

        while True:
            message = await websocket.receive_text()
            if message == "heartbeat":
                await manager.send_personal_message(message, websocket)
                continue

            # await manager.send_personal_message("系统维护中，暂停使用", websocket)
            # continue

            dictData = json.loads(message)
            print(message)
            prompt = dictData["prompt"]
            # flag备用
            flag = dictData["flag"]

            resp = "No message"
            voiceText = resp
            if flag == 0:

                with get_openai_callback() as cb:
                    resp = agent.run(prompt)
                    print(f'Total tokens: {cb.total_tokens}')
                    print(f'Requests: {cb.successful_requests}')
                voiceText = resp
            elif flag == 1:
                newsData = getGoogleNewsData(prompt)
                resp = getGoogleNewsHtml(newsData)
                voiceText = getGoogleNewsSummary(newsData)
            else:
                resp = "No message"

            timestamp = str(int(time.time()))

            subprocess.Popen(['python3', 'text2speech.py', voiceText,
                             f'./static/audio/{timestamp}.mp3'], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
            await manager.send_personal_message(f"{resp}>*_*<{timestamp}", websocket)

    except Exception as e:
        print(f'error: {e}')
        manager.disconnect(websocket)

if __name__ == "__main__":
    uvicorn.run(host="0.0.0.0", port=8000, app=app)
